// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef MODEL_EVALUATOR_FACTORY_HPP
#define MODEL_EVALUATOR_FACTORY_HPP

#include <iostream>
#include <string>
#include <map>

#include "FieldManagerBuilder.hpp"

#include "Teuchos_RCP.hpp"
#include "Teuchos_Ptr.hpp"
#include "Teuchos_Comm.hpp"
#include "Teuchos_DefaultMpiComm.hpp"
#include "Teuchos_ParameterList.hpp"
#include "Teuchos_ParameterListAcceptorDefaultBase.hpp"
#include "Tpetra_Core.hpp"

#include "PanzerAdaptersSTK_config.hpp"
#include "Panzer_STK_MeshFactory.hpp"
#include "Panzer_STK_Interface.hpp"
#include "Panzer_IntrepidFieldPattern.hpp"
#include "Panzer_ResponseLibrary.hpp"
#include "Panzer_EquationSet_Factory.hpp"
#include "Panzer_BCStrategy_Factory.hpp"
#include "Panzer_ClosureModel_Factory_TemplateManager.hpp"
#include "Panzer_ModelEvaluator.hpp"
#include "Panzer_NodeType.hpp"

#include "Piro_PerformSolve.hpp"

#include "Teko_RequestHandler.hpp"
//#include "TempusObserverFactory.hpp"
//#include "Tempus_IntegratorObserverBasic.hpp"
//#include "Tempus_IntegratorBasic.hpp"

#include "Thyra_ModelEvaluator.hpp"

namespace Piro {
  template <typename ScalarT> class RythmosSolver;
  template <typename ScalarT> class TempusSolver;
}

namespace Thyra {
  template<typename ScalarT> class ModelEvaluator;
  template<typename ScalarT> class LinearOpWithSolveFactoryBase;
}

namespace panzer {
  struct GlobalData;
  class GlobalIndexer;
  template <typename> class LinearObjFactory;

  class BlockedDOFManager;
  class DOFManager;
  class ConnManager;
}

namespace panzer_stk {
  class STKConnManager;
  class NOXObserverFactory;
  class WorksetFactory;
}

namespace SiYuan {
  class TempusObserverFactory;
}

namespace SiYuan {

  template<typename ScalarT>
  class ModelEvaluatorFactory : public Teuchos::ParameterListAcceptorDefaultBase {

  public:

    /** @name Overridden from ParameterListAcceptor */
    //@{
    void setParameterList(Teuchos::RCP<Teuchos::ParameterList> const& paramList);
    Teuchos::RCP<const Teuchos::ParameterList> getValidParameters() const;
    //@}

    /** \brief Builds the model evaluators for a panzer assembly

        \param[in] comm (Required) Teuchos communicator.  Must be non-null.
        \param[in] global_data (Required) A fully constructed (all members allocated) global data object used to control parameter library and output support. Must be non-null.
        \param[in] eqset_factory (Required) Equation set factory to provide user defined equation sets.
        \param[in] bc_factory (Required) Boundary condition factory to provide user defined boundary conditions.
        \param[in] cm_factory (Required) Closure model factory to provide user defined closure models.
    */
    void buildObjects(const Teuchos::RCP<const Teuchos::Comm<int> >& comm,
                      const Teuchos::RCP<panzer::GlobalData>& global_data,
                      const Teuchos::RCP<const panzer::EquationSetFactory>& eqset_factory,
                      const panzer::BCStrategyFactory & bc_factory,
                      const panzer::ClosureModelFactory_TemplateManager<panzer::Traits> & cm_factory,
                      bool meConstructionOn=true);

    void setInitialConditions(Teuchos::ParameterList&, 
                      Thyra::ModelEvaluatorDefaultBase<ScalarT> & model,
                      Teuchos::RCP<panzer::GlobalIndexer>& globalIndexer,
	                    Teuchos::RCP<panzer::LinearObjFactory<panzer::Traits> >& lof,
                      Teuchos::RCP<panzer::PhysicsBlock>& physicsBlock);

    Teuchos::RCP<Thyra::ModelEvaluator<ScalarT> > getPhysicsModelEvaluator();

    /** @name Methods for building the solver */
    //@{
    template <typename BuilderT>
    int addResponse(const std::string & responseName,const std::vector<panzer::WorksetDescriptor> & wkstDesc,const BuilderT & builder);

    void buildResponses(const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& cm_factory,
                        const bool write_graphviz_file=false,
                        const std::string& graphviz_file_prefix="");
    void buildIOResponses(const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& cm_factory,
                        Teuchos::RCP<Teuchos::ParameterList>& input_params,
                        const bool write_graphviz_file=false,
                        const std::string& graphviz_file_prefix="");
    void buildFluxResponses(const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& cm_factory,
                        Teuchos::RCP<Teuchos::ParameterList>& input_params,
                        const bool write_graphviz_file=false,
                        const std::string& graphviz_file_prefix="");
    void evalResponses(Teuchos::RCP<Thyra::VectorBase<double> >& gx);

    Teuchos::RCP<Thyra::ModelEvaluator<ScalarT> > getResponseOnlyModelEvaluator();

    Teuchos::RCP<Thyra::ModelEvaluator<ScalarT> >
    buildControlModelEvaluator(const Teuchos::RCP<Thyra::ModelEvaluator<ScalarT> > & thyra_me,
                                    const Teuchos::RCP<panzer::GlobalData>& global_data );

    //@}

    //! Set user defined workset factory
    void setUserWorksetFactory(Teuchos::RCP<panzer_stk::WorksetFactory>& user_wkst_factory);

    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > getResponseLibrary()
    {
      TEUCHOS_TEST_FOR_EXCEPTION(Teuchos::is_null(m_response_library), std::runtime_error,
                       "Objects are not built yet!  Please call buildObjects() member function.");

      return m_response_library;
    }

    const std::vector<Teuchos::RCP<panzer::PhysicsBlock> > & getPhysicsBlocks() const
    {
      TEUCHOS_TEST_FOR_EXCEPTION(m_physics_blocks.size()==0, std::runtime_error,
                       "Objects are not built yet!  Please call buildObjects() member function.");

      return m_physics_blocks;
    }

    //! Get mesh object used to build model evaluator
    Teuchos::RCP<panzer_stk::STK_Interface> getMesh() const
    { return m_mesh; }

    //! Get global indexer used to build model evaluator
    Teuchos::RCP<panzer::GlobalIndexer> getGlobalIndexer() const
    { return m_global_indexer; }

    //! Get connection manager
    Teuchos::RCP<panzer::ConnManager> getConnManager() const
    { return m_conn_manager; }

    //! Is blocked assembly?
    bool isBlockedAssembly() const
    { return m_blockedAssembly; }

    //! Get linear object factory used to build model evaluator
    Teuchos::RCP<panzer::LinearObjFactory<panzer::Traits> > getLinearObjFactory() const
    { return m_lin_obj_factory; }

    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > getSTKIOResponseLibrary() const
    { return m_stkIOResponseLibrary; }

    bool isTransient() const
    { return m_is_transient; }

    bool isDotDot() 
    { return m_is_dotdot; }

    void enableDotDot()
    { m_is_dotdot = true; }


    /** \brief Write the initial conditions to exodus. Note that this
      *        is entirely self contained.
      */
    void writeInitialConditions(const Thyra::ModelEvaluator<ScalarT> & model,
                                const std::vector<Teuchos::RCP<panzer::PhysicsBlock> >& physicsBlocks,
                                const Teuchos::RCP<panzer::WorksetContainer> & wc,
                                const Teuchos::RCP<const panzer::GlobalIndexer> & ugi,
                                const Teuchos::RCP<const panzer::LinearObjFactory<panzer::Traits> > & lof,
                                const Teuchos::RCP<panzer_stk::STK_Interface> & mesh,
                                const panzer::ClosureModelFactory_TemplateManager<panzer::Traits> & cm_factory,
                                const Teuchos::ParameterList & closure_model_pl,
                                const Teuchos::ParameterList & user_data_pl,
                                int workset_size) const;

    bool useDynamicCoordinates() const
    { return useDynamicCoordinates_; }

    Teuchos::RCP<Thyra::LinearOpWithSolveFactoryBase<double> >
    buildLOWSFactory(bool blockedAssembly,
                     const Teuchos::RCP<const panzer::GlobalIndexer> & globalIndexer,
                     const Teuchos::RCP<panzer::ConnManager> & conn_manager,
                     const Teuchos::RCP<panzer_stk::STK_Interface> & mesh,
                     const Teuchos::RCP<const Teuchos::MpiComm<int> > & mpi_comm,
                     const Teuchos::RCP<Teko::RequestHandler> & req_handler=Teuchos::null
                     ) const;

    //! Get the workset container associated with the mesh database.
    Teuchos::RCP<panzer::WorksetContainer> getWorksetContainer() const
    { return m_wkstContainer; }

    //! Add the user fields specified by output_list to the mesh
    void addUserFieldsToMesh(panzer_stk::STK_Interface & mesh,const Teuchos::ParameterList & output_list) const;

    //! build STK mesh factory from a mesh parameter list
    Teuchos::RCP<panzer_stk::STK_MeshFactory> buildSTKMeshFactory(const Teuchos::ParameterList & mesh_params) const;

    void finalizeMeshConstruction(const panzer_stk::STK_MeshFactory & mesh_factory,
                                  const std::vector<Teuchos::RCP<panzer::PhysicsBlock> > & physicsBlocks,
                                  const Teuchos::MpiComm<int> mpi_comm,
                                  panzer_stk::STK_Interface & mesh) const;

  protected:

    Teuchos::RCP<FieldManagerBuilder>
    buildFieldManagerBuilder(const Teuchos::RCP<panzer::WorksetContainer> & wc,
                             const std::vector<Teuchos::RCP<panzer::PhysicsBlock> >& physicsBlocks,
                             const std::vector<panzer::BC> & bcs,
                             const std::vector< Dirichlet<panzer::Traits::RealType> >& ds,
                             const std::vector< NeumannBoundary >& ns,
                             const std::vector< CLoad >& cloads,
                             const panzer::EquationSetFactory & eqset_factory,
                             const panzer::BCStrategyFactory& bc_factory,
                             const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& volume_cm_factory,
                             const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& bc_cm_factory,
                             const Teuchos::ParameterList& closure_models,
                             const panzer::LinearObjFactory<panzer::Traits> & lo_factory,
                             const Teuchos::ParameterList& user_data,
                             bool writeGraph,const std::string & graphPrefix,
			     bool write_field_managers,const std::string & field_manager_prefix) const;

    /**
      */
    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > initializeSolnWriterResponseLibrary(
                                                                const Teuchos::RCP<panzer::WorksetContainer> & wc,
                                                                const Teuchos::RCP<const panzer::GlobalIndexer> & ugi,
                                                                const Teuchos::RCP<const panzer::LinearObjFactory<panzer::Traits> > & lof,
                                                                const Teuchos::RCP<panzer_stk::STK_Interface> & mesh) const;

    /**
      */
    void finalizeSolnWriterResponseLibrary(panzer::ResponseLibrary<panzer::Traits> & rl,
                                           const std::vector<Teuchos::RCP<panzer::PhysicsBlock> > & physicsBlocks,
                                           const panzer::ClosureModelFactory_TemplateManager<panzer::Traits> & cm_factory,
                                           const Teuchos::ParameterList & closure_models,
                                           int workset_size, Teuchos::ParameterList & user_data) const;


  private:

    Teuchos::RCP<Thyra::ModelEvaluator<ScalarT> > m_physics_me;
    Teuchos::RCP<Thyra::ModelEvaluator<ScalarT> > m_rome_me;

    //volume response
    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > m_response_library;
    std::map<int,std::string> m_responseIndexToName;
    // surface reponse
    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > m_fluxResponse;
    std::vector<Teuchos::RCP<panzer::PhysicsBlock> > m_physics_blocks;

    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > m_stkIOResponseLibrary;

    Teuchos::RCP<panzer_stk::STK_Interface> m_mesh;
    Teuchos::RCP<panzer::GlobalIndexer> m_global_indexer;
    Teuchos::RCP<panzer::ConnManager> m_conn_manager;
    Teuchos::RCP<panzer::LinearObjFactory<panzer::Traits> > m_lin_obj_factory;
    Teuchos::RCP<panzer::GlobalData> m_global_data;
    bool useDiscreteAdjoint;
    bool m_is_transient;
    bool m_is_dotdot;
    bool m_blockedAssembly;
    Teuchos::RCP<const panzer::EquationSetFactory> m_eqset_factory;

    Teuchos::RCP<panzer_stk::WorksetFactory> m_user_wkst_factory;
    Teuchos::RCP<panzer::WorksetContainer> m_wkstContainer;

    bool useDynamicCoordinates_;
  };

template<typename ScalarT>
template <typename BuilderT>
int ModelEvaluatorFactory<ScalarT>::
addResponse(const std::string & responseName,const std::vector<panzer::WorksetDescriptor> & wkstDesc,const BuilderT & builder)
{
  typedef panzer::ModelEvaluator<double> PanzerME;
  Teuchos::RCP<PanzerME> panzer_me = Teuchos::rcp_dynamic_cast<PanzerME>(m_physics_me);

  if(panzer_me!=Teuchos::null) {
    return panzer_me->addResponse(responseName,wkstDesc,builder);
  }

  TEUCHOS_ASSERT(false);
  return -1;
}

}

#include "ModelEvaluatorFactory_impl.hpp"

#endif
